import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import math
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import os
from matplotlib.colors import LogNorm

# 模型定义
def create_linear_model():
    model = nn.Sequential(
        nn.Linear(1, 32),
        nn.ReLU(),
        nn.Linear(32, 128),
        nn.ReLU(),
        nn.Linear(128, 64),
        nn.ReLU(),
        nn.Linear(64, 1)
    )
    return model

def data_Min_Max():
    # 读取原始数据
    input_path = os.path.join('数据', 'CS2_result.xlsx')
    train_Sheet = 'Sheet1'
    dataset = pd.read_excel(input_path, sheet_name=train_Sheet)

    x = dataset.iloc[:, 0].values.reshape(-1, 1).astype(float)
    y = dataset.iloc[:, 1].values.reshape(-1, 1).astype(float)

    # 打乱数据 (在原始数据上打乱)
    np.random.seed(5)
    indices = np.random.permutation(len(x))
    x_shuffled = x[indices]
    y_shuffled = y[indices]

    # 划分训练测试集 (使用原始数据)
    split_index = int(len(x_shuffled) * 0.8)
    train_x_raw, test_x_raw = x_shuffled[:split_index], x_shuffled[split_index:]
    train_y_raw, test_y_raw = y_shuffled[:split_index], y_shuffled[split_index:]

    # 计算归一化参数 (仅用训练集)
    min_x, max_x = np.min(train_x_raw), np.max(train_x_raw)
    min_y, max_y = np.min(train_y_raw), np.max(train_y_raw)

    # 避免除零错误 (如果数据全为同一值)
    range_x = max_x - min_x if (max_x - min_x) != 0 else 1.0
    range_y = max_y - min_y if (max_y - min_y) != 0 else 1.0

    # 执行Min-Max归一化
    train_x = (train_x_raw - min_x) / range_x
    train_y = (train_y_raw - min_y) / range_y
    test_x = (test_x_raw - min_x) / range_x
    test_y = (test_y_raw - min_y) / range_y

    return train_x, train_y, test_x, test_y, min_y, max_y
# 莱维飞行生成随机步长
def levy_flight(dim):
    beta = 1.2
    sigma = (math.gamma(1 + beta) * np.sin(np.pi * beta / 2) / (
                math.gamma((1 + beta) / 2) * beta * 2 ** ((beta - 1) / 2))) ** (1 / beta)
    u = np.random.normal(0, sigma, dim)
    v = np.random.normal(0, 1, dim)
    step = u / (np.abs(v) ** (1 / beta))
    return step


# 梯度下降局部优化（核心改进）
def gradient_descent_step(model, weights, train_loader, lr=0.01, steps=10):
    """对当前权重进行多步梯度下降优化"""
    set_weight(model, weights)  # 加载权重到模型

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    #执行多步梯度下降
    for _ in range(steps):
        model.train()
        for inputs, targets in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_function(outputs, targets)
            loss.backward()
            optimizer.step()

    # 将更新后的权重扁平化返回
    new_weights = np.concatenate([param.detach().cpu().numpy().flatten() for param in model.parameters()])
    return new_weights


# 布谷鸟搜索算法（含动态步长和局部优化）
def cuckoo_search(model, train_loader, loss_func, n_nest=100, pa=0.3, max_iter=3000):
    # 初始化种群
    init_weights = [p.detach().cpu().numpy() for p in model.parameters()]
    dim = sum(w.size for w in init_weights)
    #初始参数一维
    init_flat = np.concatenate([w.flatten() for w in init_weights])
    #生成输出鸟巢
    nests = np.array([init_flat + np.random.randn(dim) * 0.01 for _ in range(n_nest)])
    best_nest = nests[0].copy()
    #初始损失
    best_loss = evaluate_model(model, train_loader, loss_func, best_nest)

    for epoch in range(max_iter):
        # 动态参数计算（步长和学习率随迭代衰减）
        decay = np.exp(-epoch / (max_iter * 0.5))
        lr = 0.001 * decay

        # ---------- 全局搜索 ----------
        for i in range(n_nest):
            # 莱维飞行公式
            step = 0.1 * decay * levy_flight(dim)
            new_nest = nests[i] + step * np.random.randn(dim)

            # 评估与更新
            curr_loss = evaluate_model(model, train_loader, loss_func, new_nest)
            if curr_loss < best_loss:
                best_nest = new_nest.copy()
                best_loss = curr_loss
                nests[i] = new_nest.copy()

        # ---------- 局部搜索 ----------
        for j in range(n_nest):
            #按照概率进行局部优化
            if np.random.rand() < pa:
                # 差分扰动公式
                r1, r2 = np.random.choice([k for k in range(n_nest) if k != j], 2, False)
                e = np.random.rand()
                H = 1 if (e - pa) > 0 else 0
                # 修改后（动态扰动因子，随迭代衰减）
                perturbation_factor = 0.5 * decay  # 扰动强度随迭代衰减
                perturbation = perturbation_factor * H * (nests[r1] - nests[r2])

                # 生成新解
                new_nest = nests[j] + perturbation
                new_nest = gradient_descent_step(model, new_nest, train_loader, lr=lr)

                # 择优替换
                new_loss = evaluate_model(model, train_loader, loss_func, new_nest)
                old_loss = evaluate_model(model, train_loader, loss_func, nests[j])
                if new_loss < old_loss:
                    nests[j] = new_nest.copy()
                    if new_loss < best_loss:
                        best_nest = new_nest.copy()
                        best_loss = new_loss

        # 早停检测
        if epoch > 500 and abs(prev_loss - best_loss) < 1e-6:
            break
        prev_loss = best_loss
        if (epoch) % 100 == 0:
            print("Epoch,loss",epoch,best_loss)

    return best_nest


# 其他辅助函数（保持不变）
def evaluate_model(model, train_loader, loss_function, weights):
    set_weight(model, weights)
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for inputs, targets in train_loader:
            outputs = model(inputs)
            loss = loss_function(outputs, targets)
            total_loss += loss.item()
    return total_loss / len(train_loader)


def set_weight(model, weights):
    idx = 0
    for param in model.parameters():
        param_size = param.numel()
        param_shape = param.shape
        param.data = torch.tensor(weights[idx:idx + param_size].reshape(param_shape), dtype=torch.float32)
        idx += param_size

# 设置全局绘图参数（符合SCI期刊格式要求）
plt.rcParams.update({
    "font.family": "serif",  # 使用衬线字体
    "font.serif": ["Times New Roman"],  # 指定Times New Roman字体
    "font.size": 11,  # 基础字体大小
    "axes.labelsize": 11,  # 坐标轴标签字体大小
    "axes.titlesize": 12,  # 标题字体大小
    "xtick.labelsize": 10,  # X轴刻度标签大小
    "ytick.labelsize": 10,  # Y轴刻度标签大小
    "xtick.direction": "in",  # 刻度线朝内
    "ytick.direction": "in",  # 刻度线朝内
    "axes.linewidth": 1.5,  # 坐标轴线宽
    "lines.linewidth": 1.5,  # 绘图线宽
    "savefig.dpi": 600,  # 保存图像分辨率
    "mathtext.fontset": "stix",  # 数学公式字体与正文一致
})

def plot_sci_scatter(true, pred):
    fig, ax = plt.subplots(figsize=(4.5, 4.5), dpi=600)

    # ============== 计算相对误差 ==============
    epsilon = 1e-8  # 防止除以0
    relative_error = np.abs((pred - true) / (true + epsilon))  # 相对误差绝对值

    # ============== 绘制散点（颜色映射改为相对误差） ==============
    sc = ax.scatter(
        true, pred,
        c=relative_error,  # 颜色映射基于相对误差
        s=35,
        cmap='coolwarm',  # 改为单色渐变色
        alpha=0.8,
        edgecolors='w',
        linewidths=0.4,
        norm=LogNorm(vmin=1e-4, vmax=1),  # 对数归一化（适用于误差量级差异大的情况）
        zorder=3
    )

    # ============== 参考线配置 ==============
    # 智能坐标范围（排除1%极端值，保证主要数据区域显示）
    data_min = np.percentile([true, pred], 1)  # 1%分位数
    data_max = np.percentile([true, pred], 99)  # 99%分位数
    ax.plot(
        [data_min, data_max], [data_min, data_max],
        '--',  # 虚线样式
        color='#2d2d2d',  # 深灰色提高对比度
        lw=1.2,  # 线宽略粗于默认值
        alpha=0.9,  # 适当降低透明度
        zorder=2  # 绘制层级（在散点下方）
    )

    # ============== 颜色条优化 ==============
    cbar = fig.colorbar(sc,ax=ax,pad=0.02,aspect=30)
    cbar.ax.tick_params(labelsize=9)  # 颜色条刻度标签大小
    cbar.set_label('Relative Error', fontsize=10, labelpad=2)

    # ============== 坐标轴与标签 ==============
    ax.set(
        xlim=(data_min, data_max),  # X轴范围与参考线对齐
        ylim=(data_min, data_max),  # Y轴范围相同保证正方形显示
        xlabel='True Value',  # X轴标签（根据实际数据替换）
        ylabel='Predicted Value'  # Y轴标签
    )

    # ============== 统计信息标注 ==============
    stats_text = ('(a)')
    ax.text(
        0.05, 0.95,  # 位于左上角（坐标系相对位置）
        stats_text,
        transform=ax.transAxes,  # 使用相对坐标系
        ha='left', va='top',  # 左对齐，顶部对齐
        fontsize=10,
        bbox=dict(
            facecolor='white', alpha=0.9,  # 半透明白底背景
            edgecolor='0.8',  # 浅灰色边框
            boxstyle='round,pad=0.3'  # 圆角边框，内边距0.3
        )
    )

    # ============== 网格与边框优化 ==============
    ax.grid(True,ls=':',color='gray',alpha=0.4)
    # 强化边框线宽（与刻度线协调）
    for spine in ax.spines.values():
        spine.set_linewidth(1.2)

    # ============== 布局与保存 ==============
    plt.tight_layout(pad=1.5)  # 紧凑布局（留白控制）

    # 创建输出目录（如果不存在）
    os.makedirs('PDF', exist_ok=True)

    # 保存双版本图片（位图+矢量图）
    save_path = os.path.join('PDF', 'F_cs2')
    plt.savefig(f'{save_path}.pdf', transparent=True)  # 矢量图适合编辑
    plt.close()  # 显式关闭图形释放内存

# 测试函数
def test(model,test_loader,mean_y, std_y):
    model.eval()
    pre_outputs = []  # 存储预测值
    true_outputs = []  # 存储实际值
    with torch.no_grad():
        for input, target in test_loader:
            output = model(input)
            # 将预测值转换为numpy数组
            pre_output_np = output.cpu().numpy()
            pre_outputs.append(pre_output_np * std_y + mean_y)
            # pre_outputs.append(pre_output_np)

            # 将实际值转换为numpy数组并反归一化
            true_output_np = target.cpu().numpy()
            true_outputs.append(true_output_np * std_y + mean_y)
            # true_outputs.append(true_output_np )

    # 将列表转换为numpy数组，并合并为一个二维数组
    pre_outputs = np.concatenate(pre_outputs, axis=0).flatten()
    true_outputs = np.concatenate(true_outputs, axis=0).flatten()

    mse = mean_squared_error(true_outputs, pre_outputs)
    rmse = np.sqrt(mse)

    mae = mean_absolute_error(true_outputs, pre_outputs)
    r2 = r2_score(true_outputs, pre_outputs)
    mape = np.mean(np.abs((true_outputs - pre_outputs) / true_outputs) * 100)

    print(f"MSE: {mse:.4f}, RMSE: {rmse:.4f}")
    print(f"MAE: {mae:.4f}, R²: {r2:.4f},MAPE:{mape:.4f}")

    # 保存结果
    df = pd.DataFrame({
        'Actual': true_outputs,
        'Predicted': pre_outputs,
        'Error': true_outputs - pre_outputs,
        'Relative Error': (true_outputs - pre_outputs) / true_outputs
    })
    os.makedirs('结果', exist_ok=True)
    df.to_excel(os.path.join('结果', 'F_CS2_Results.xlsx'), index=False)

    # 调用SCI绘图函数
    plot_sci_scatter(true_outputs, pre_outputs)

if __name__=="__main__":
    train_inputs, train_outputs, test_inputs, test_outputs, mean_y, std_y = data_Min_Max()
    train_dataset = TensorDataset(torch.tensor(train_inputs, dtype=torch.float32),
                                  torch.tensor(train_outputs, dtype=torch.float32))
    test_dataset = TensorDataset(torch.tensor(test_inputs, dtype=torch.float32),
                                 torch.tensor(test_outputs, dtype=torch.float32))
    BATCH_SIZE = 32
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
    model = create_linear_model()
    loss_function = nn.MSELoss(reduction='mean')
    best_weights = cuckoo_search(model, train_loader, loss_function)
    set_weight(model, best_weights)

    model_path = os.path.join('模型', "F_CS2_model.pth")

    # 保存模型权重
    torch.save(model.state_dict(), model_path)
    model_path = os.path.join('模型', "F_CS2_model.pth")
    model.load_state_dict(torch.load(model_path))
    test(model, test_loader,mean_y, std_y)